import mindspore as ms
from mindspore import nn
from mindspore import ops
from mindspore.dataset import vision, transforms
from mindspore.dataset import MnistDataset
from mindspore.train import Model, CheckpointConfig, ModelCheckpoint, LossMonitor

from wpgpmfm.unet3d.networks.unet3d import Unet3d
from wpgpmfm.unet3d.datasets.loaders import HematomaDataset

ROOT = '/share_data/liupan/hematoma/npyfiles/24h/train'
# files = [os.path.join(ROOT, f) for f in os.listdir(ROOT)]
train_ds = HematomaDataset(ROOT)
train_dsl = ms.dataset.GeneratorDataset(train_ds, column_names=['img', 'lab'], shuffle=True)
train_dsl = train_dsl.batch(2, drop_remainder=True)
train_dsl = train_dsl.repeat(1)

model = Unet3d(1, 2)
loss_fn = nn.CrossEntropyLoss()
optimizer = nn.SGD(model.trainable_params(), 1e-3)

steps_per_epoch = train_dsl.get_dataset_size()
config = CheckpointConfig(save_checkpoint_steps=steps_per_epoch)

ckpt_callback = ModelCheckpoint(prefix="hematoma", directory="./checkpoint", config=config)
loss_callback = LossMonitor(steps_per_epoch)

trainer = Model(model, loss_fn=loss_fn, optimizer=optimizer, metrics={'accuracy'})

trainer.train(10, train_dsl, callbacks=[ckpt_callback, loss_callback])
